/*
 * @Author: Sian Li lisian_magic@163.com
 * @Date: 2024-08-21 16:44:50
 * @LastEditors: Sian Li lisian_magic@163.com
 * @LastEditTime: 2024-08-21 18:22:16
 * @FilePath: /Yolov5Detect/Detectors/include/ArmorDetector.h
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置
 * 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */
#ifndef __DETECTOR_H__
#define __DETECTOR_H__
#include "../TRTInference/TRTInfer/include/Inference.h"
#include <cstdio>
#include <fcntl.h>
#include <unistd.h>

using namespace cv;
using namespace std;
/**
 * @brief 一层网络检测框
 */
typedef struct DetectBox {
    DetectBox(float x1 = 0, float y1 = 0, float x2 = 0, float y2 = 0,
              float confidence = 0, float classID = -1, float trackID = -1) {
        this->x1 = x1;
        this->y1 = y1;
        this->x2 = x2;
        this->y2 = y2;
        this->confidence = confidence;
        this->classID = classID;
        this->trackID = trackID;
    }
    float x1, y1, x2, y2;
    float confidence;
    float classID;
    float trackID;
} DetectBox;

/**
 * @brief 装甲板BBox
 */
struct ArmorBoundingBox {
    bool flag = false;
    _Float32 x0 = 0.f, y0 = 0.f, w = 0.f, h = 0.f, cls = 0.f, conf = 0.f,
             depth = 0.f;
};

/**
 * @brief 装甲板BBox和分割Rect绑定包
 */
struct bboxAndRect {
    ArmorBoundingBox armor;
    DetectBox rect;
};
/**
 * @brief 装甲板识别
 * 装甲板识别器，输出装甲板BBox和对应分割框
 */
class ArmorDetector {
  public:
    typedef std::shared_ptr<ArmorDetector> Ptr;

  private:
    TRTInferV1::TRTInfer armorTensorRT = TRTInferV1::TRTInfer(0);
    vector<bboxAndRect> results;

    string TensorRTEnginePath;
    string OnnxPath;

  private:
    vector<Mat> preProcess(Mat &image, vector<DetectBox> &movingTargets);
    void reBuildBoxs(vector<vector<TRTInferV1::DetectionObj>> &armors,
                     vector<DetectBox> &boxs, vector<Mat> &img);

  public:
    ArmorDetector(string engine_path, string onnx_path);
    ~ArmorDetector();

    void accessModelTest();
    bool initModel();
    vector<bboxAndRect> infer(Mat &image, vector<DetectBox> &targets);
    void unInit();
};

class CarDetector {
  public:
    typedef std::shared_ptr<CarDetector> Ptr;

  private:
    TRTInferV1::TRTInfer carTensorRT = TRTInferV1::TRTInfer(0);

    string TensorRTEnginePath;
    string OnnxPath;

  public:
    CarDetector(string engine_path, string onnx_path);
    ~CarDetector();

    void accessModelTest();
    bool initModel();
    vector<DetectBox> infer(Mat &image);
    void unInit();
};

#endif